import numpy as np
import torch
from imageio import imread
from matplotlib import pyplot as plt
from skimage.transform import resize as imresize

import models

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
pretrained_dispnet = "/media/gzy/My Passport/Unsupervised_Monocular_Depth_Estimation/checkpoint/resnet_laplace/09-27-15:26/dispnet_model_best.pth.tar"
dispnet = "DispResNet"
disp_net = getattr(models, dispnet)(16).to(device) if dispnet in ['DispSENet', 'DispDenseNet'] else getattr(models, dispnet)().to( device)

weights = torch.load(pretrained_dispnet, map_location=torch.device('cpu'))
disp_net.load_state_dict(weights['state_dict'])
disp_net.eval()


def load_tensor_image(filename, img_height, img_width):
    img = imread(filename).astype(np.float32)
    h, w, _ = img.shape
    if h != img_height or w != img_width:
        img = imresize(img, (img_height, img_width)).astype(np.float32)
    img = np.transpose(img, (2, 0, 1))
    tensor_img = ((torch.from_numpy(img).unsqueeze(0) / 255 - 0.5) / 0.5).to(device)
    return tensor_img


def test_image():
    test_file = "/media/gzy/My Passport/Unsupervised_Monocular_Depth_Estimation/kitti_256/2011_09_26_drive_0057_sync_03/0000000290.jpg"
    tgt_img = load_tensor_image(test_file, 256, 832)
    print(tgt_img.shape)
    pred_disp = disp_net(tgt_img).cpu().detach().numpy()[0, 0]
    predictions = 1 / pred_disp
    print(type(pred_disp))
    print(pred_disp.shape)
    plt.imshow(pred_disp, vmin=0, vmax=pred_disp.max(), cmap='turbo')
    plt.show()
    pass


def main():
    test_image()
    # print(disp_net.to(device))
    pass


if __name__ == '__main__':
    main()
